import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import threading
import os
from dotenv import load_dotenv
import base64
import urllib.parse
import requests
import pandas as pd
import json
import openpyxl


# API_KEY = "8zMXXXXXXGb2cI"
# SECRET_KEY = "sHo0XXXXXnuXXXXXXjkR"
# 读取env文件中的API_KEY和SECRET_KEY
load_dotenv()
API_KEY = os.getenv("API_KEY")
SECRET_KEY = os.getenv("SECRET_KEY")

# 字段名称映射
INVOICE_FIELD_MAPPING = {
    "log_id": "日志ID",
    "pdf_file_size": "PDF文件页数",
    "ofd_file_size": "OFD文件页数",
    "words_result_num": "识别结果数",
    "ServiceType": "服务类型",
    "InvoiceType": "发票种类",
    "InvoiceTypeOrg": "发票名称",
    "InvoiceCode": "发票代码",
    "InvoiceNum": "发票号码",
    "InvoiceCodeConfirm": "发票代码校验码",
    "InvoiceNumConfirm": "发票号码校验码",
    "InvoiceNumDigit": "数电票号",
    "InvoiceTag": "发票标志",
    "MachineNum": "机打号码",
    "MachineCode": "机器编号",
    "CheckCode": "校验码",
    "InvoiceDate": "开票日期",
    "PurchaserName": "购方名称",
    "PurchaserRegisterNum": "购方纳税人识别号",
    "PurchaserAddress": "购方地址及电话",
    "PurchaserBank": "购方开户行及账号",
    "Password": "密码区",
    "Province": "省",
    "City": "市",
    "SheetNum": "联次信息",
    "Agent": "是否代开",
    "SellerName": "销售方名称",
    "SellerRegisterNum": "销售方纳税人识别号",
    "SellerAddress": "销售方地址及电话",
    "SellerBank": "销售方开户行及账号",
    "TotalAmount": "合计金额",
    "TotalTax": "合计税额",
    "AmountInWords": "价税合计(大写)",
    "AmountInFiguers": "价税合计(小写)",
    "Payee": "收款人",
    "Checker": "复核",
    "NoteDrawer": "开票人",
    "Remarks": "备注",
    "company_seal": "公司印章",
    "seal_info": "印章信息"
}

PRODUCT_FIELD_MAPPING = {
    "行号": "行号",
    "CommodityName": "货物名称",
    "CommodityType": "规格型号",
    "CommodityNum": "数量",
    "CommodityAmount": "金额",
    "CommodityPrice": "单价",
    "CommodityTaxRate": "税率",
    "CommodityUnit": "单位",
    "CommodityTax": "税额"
}

def main(file_paths, output_filename='output.xlsx'):
    invoice_records = []
    product_records = []

    for file_path in file_paths:
        print(f"处理文件: {file_path}")
        response_json = get_ocr_result(file_path)
        if response_json:
            invoice_info, product_details = extract_data(response_json)
            if invoice_info:
                invoice_records.append(invoice_info)
            if product_details:
                product_records.extend(product_details)
        else:
            print(f"未能获取OCR结果: {file_path}")

    if invoice_records and product_records:
        save_to_excel(invoice_records, product_records, output_filename)
        print(f"所有数据已成功保存到 {output_filename}")
    else:
        print("没有足够的数据保存到Excel。")

def get_ocr_result(file_path):
    access_token = get_access_token()
    if not access_token:
        print("无法获取Access Token，跳过文件。")
        return None

    url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/vat_invoice?access_token={access_token}"

    payload = 'pdf_file=' + get_file_content_as_base64(file_path, True) + '&seal_tag=false'
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded',
        'Accept': 'application/json'
    }

    try:
        response = requests.post(url, headers=headers, data=payload)
        if response.status_code == 200:
            try:
                return response.json()
            except json.JSONDecodeError:
                print(f"文件 {file_path} 的响应内容不是有效的JSON格式")
                return None
        else:
            print(f"请求失败，状态码: {response.status_code}，文件: {file_path}")
            return None
    except requests.RequestException as e:
        print(f"请求异常: {e}，文件: {file_path}")
        return None

def extract_data(data):
    words = data.get("words_result", {})

    # 如果 words_result 是字符串，则尝试解析为字典
    if isinstance(words, str):
        try:
            words = json.loads(words)
        except json.JSONDecodeError:
            print("words_result 不是有效的JSON字符串")
            words = {}

    # 确保 words 是字典
    if not isinstance(words, dict):
        print("words_result 的格式不正确")
        return None, None

    # 发票信息提取
    invoice_info = {}
    for field, chinese_name in INVOICE_FIELD_MAPPING.items():
        field_data = words.get(field, {})
        if isinstance(field_data, dict):
            invoice_info[chinese_name] = field_data.get("word", "")
        elif isinstance(field_data, str):
            invoice_info[chinese_name] = field_data
        else:
            invoice_info[chinese_name] = ""

    # 商品详情提取
    commodity_fields = ["CommodityNum", "CommodityAmount", "CommodityPrice", 
                        "CommodityTaxRate", "CommodityUnit", "CommodityTax", 
                        "CommodityName", "CommodityType"]

    # 检查所有商品字段是否为列表
    for field in commodity_fields:
        if field in words and not isinstance(words[field], list):
            print(f"字段 {field} 的类型不是列表，跳过此字段")
            words[field] = []

    # 确保所有商品字段长度一致
    field_lengths = [len(words.get(field, [])) for field in commodity_fields]
    if not field_lengths or min(field_lengths) == 0:
        print("没有找到任何商品信息")
        return invoice_info, []

    min_length = min(field_lengths)
    if len(set(field_lengths)) != 1:
        print("商品字段的长度不一致，可能存在数据缺失。将使用最小长度的数据进行处理。")

    product_details = []
    for i in range(min_length):
        product = {}
        try:
            # 行号
            commodity_num_entry = words.get("CommodityNum", [])[i]
            product["行号"] = commodity_num_entry.get("row", "") if isinstance(commodity_num_entry, dict) else ""

            # CommodityName
            commodity_name_entry = words.get("CommodityName", [])[i]
            if isinstance(commodity_name_entry, dict):
                name = commodity_name_entry.get("word", "").replace('\n', '').replace(' ', '')
            elif isinstance(commodity_name_entry, str):
                name = commodity_name_entry.replace('\n', '').replace(' ', '')
            else:
                name = ""
            product["货物名称"] = name

            # CommodityType
            commodity_type_entry = words.get("CommodityType", [])[i]
            product["规格型号"] = commodity_type_entry.get("word", "") if isinstance(commodity_type_entry, dict) else ""

            # CommodityNum
            commodity_num = commodity_num_entry.get("word", "") if isinstance(commodity_num_entry, dict) else ""
            product["数量"] = commodity_num

            # CommodityAmount
            commodity_amount_entry = words.get("CommodityAmount", [])[i]
            product["金额"] = commodity_amount_entry.get("word", "") if isinstance(commodity_amount_entry, dict) else ""

            # CommodityPrice
            commodity_price_entry = words.get("CommodityPrice", [])[i]
            product["单价"] = commodity_price_entry.get("word", "") if isinstance(commodity_price_entry, dict) else ""

            # CommodityTaxRate
            commodity_tax_rate_entry = words.get("CommodityTaxRate", [])[i]
            product["税率"] = commodity_tax_rate_entry.get("word", "") if isinstance(commodity_tax_rate_entry, dict) else ""

            # CommodityUnit
            commodity_unit_entry = words.get("CommodityUnit", [])[i]
            product["单位"] = commodity_unit_entry.get("word", "") if isinstance(commodity_unit_entry, dict) else ""

            # CommodityTax
            commodity_tax_entry = words.get("CommodityTax", [])[i]
            product["税额"] = commodity_tax_entry.get("word", "") if isinstance(commodity_tax_entry, dict) else ""

            product_details.append(product)
        except (IndexError, AttributeError) as e:
            print(f"处理第 {i+1} 行商品时出错: {e}")
            continue

    return invoice_info, product_details

def save_to_excel(invoice_records, product_records, filename):
    # 创建发票信息的DataFrame
    invoice_df = pd.DataFrame(invoice_records, columns=INVOICE_FIELD_MAPPING.values())

    # 创建商品详情的DataFrame
    product_df = pd.DataFrame(product_records, columns=PRODUCT_FIELD_MAPPING.values())

    # 写入Excel，包含两个工作表
    with pd.ExcelWriter(filename, engine='openpyxl') as writer:
        invoice_df.to_excel(writer, sheet_name='发票信息', index=False)
        product_df.to_excel(writer, sheet_name='商品详情', index=False)

def get_file_content_as_base64(path, urlencoded=False):
    try:
        with open(path, "rb") as f:
            content = base64.b64encode(f.read()).decode("utf8")
            if urlencoded:
                content = urllib.parse.quote_plus(content)
        return content
    except FileNotFoundError:
        print(f"文件未找到: {path}")
        return ""

def get_access_token():
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {
        "grant_type": "client_credentials",
        "client_id": API_KEY,
        "client_secret": SECRET_KEY
    }
    try:
        response = requests.post(url, params=params)
        if response.status_code == 200:
            try:
                return response.json().get("access_token", "")
            except json.JSONDecodeError:
                print("获取Access Token时，响应不是有效的JSON格式")
                return ""
        else:
            print(f"获取Access Token失败，状态码: {response.status_code}")
            return ""
    except requests.RequestException as e:
        print(f"请求异常: {e}")
        return ""

class App(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title('发票OCR处理工具')
        self.geometry('600x200')

        self.init_ui()

    def init_ui(self):
        self.input_dir_var = tk.StringVar()
        self.output_file_var = tk.StringVar()

        frame = ttk.Frame(self)
        frame.pack(padx=10, pady=10, fill=tk.BOTH, expand=True)

        ttk.Label(frame, text="选择发票文件夹:").grid(row=0, column=0, padx=5, pady=5, sticky=tk.W)
        ttk.Entry(frame, textvariable=self.input_dir_var, width=50).grid(row=0, column=1, padx=5, pady=5, sticky=tk.W)
        ttk.Button(frame, text="浏览", command=self.select_input_dir).grid(row=0, column=2, padx=5, pady=5, sticky=tk.W)

        ttk.Label(frame, text="选择输出文件:").grid(row=1, column=0, padx=5, pady=5, sticky=tk.W)
        ttk.Entry(frame, textvariable=self.output_file_var, width=50).grid(row=1, column=1, padx=5, pady=5, sticky=tk.W)
        ttk.Button(frame, text="浏览", command=self.select_output_file).grid(row=1, column=2, padx=5, pady=5, sticky=tk.W)

        ttk.Button(frame, text="开始处理", command=self.start_process).grid(row=2, column=0, columnspan=3, padx=5, pady=10, sticky=tk.W+tk.E)

        self.status_label = ttk.Label(frame, text="")
        self.status_label.grid(row=3, column=0, columnspan=3, padx=5, pady=5, sticky=tk.W)

    def select_input_dir(self):
        directory = filedialog.askdirectory()
        if directory:
            self.input_dir_var.set(directory)

    def select_output_file(self):
        file_path = filedialog.asksaveasfilename(defaultextension=".xlsx", filetypes=[("Excel files", "*.xlsx")])
        if file_path:
            self.output_file_var.set(file_path)

    def start_process(self):
        input_directory = self.input_dir_var.get()
        output_excel = self.output_file_var.get()

        if not input_directory or not output_excel:
            messagebox.showerror('错误', '请输入文件夹路径和输出文件名')
            return

        self.status_label.config(text="正在处理...")
        self.disable_buttons()

        thread = threading.Thread(target=self.process_files, args=(input_directory, output_excel))
        thread.start()

    def process_files(self, input_directory, output_excel):
        file_paths = [os.path.join(input_directory, file) for file in os.listdir(input_directory) if file.lower().endswith('.pdf')]
        if not file_paths:
            self.update_status('文件夹中没有找到PDF文件')
            self.enable_buttons()
            return

        try:
            main(file_paths, output_excel)
            self.update_status(f'处理完成，文件已保存至 {output_excel}')
        except Exception as e:
            self.update_status(f'处理过程中发生错误: {str(e)}')
        finally:
            self.enable_buttons()

    def update_status(self, message):
        self.status_label.config(text=message)
        self.update_idletasks()

    def disable_buttons(self):
        for widget in self.winfo_children():
            if isinstance(widget, ttk.Button):
                widget.config(state=tk.DISABLED)

    def enable_buttons(self):
        for widget in self.winfo_children():
            if isinstance(widget, ttk.Button):
                widget.config(state=tk.NORMAL)

if __name__ == '__main__':
    app = App()
    app.mainloop()